# 将数据集(kitti\4seasons\urban)的ground truth文件转换为nav_msgs.msg.Path， 放到rosbag内

import numpy as np
import os
import rospy
import rosbag
from geometry_msgs.msg import PoseStamped
from nav_msgs.msg import Path

from scipy.spatial.transform import Rotation


def get_times(filename):
    times = []
    with open(filename) as f:
        list_file = f.readlines()
        # 将每一行数据转为数组
        for i in range(len(list_file)):
            list_line = list_file[i].split(' ')
            # 将元素由字符串转为float
            list_line = list(map(float, list_line))
            times.append(list_line[0])
    return times


def write_bag_kitti(filename, bagname, times, path):
    trans = []
    quats = []
    with open(filename) as f:
        list_file = f.readlines()
        # 将每一行数据转为数组
        for i in range(len(list_file)):
            list_line = list_file[i].split()
            # 将元素由字符串转为float
            list_line = list(map(float, list_line))
            # 向量转矩阵
            list_line = np.array(list_line)
            list_line.resize(3, 4)
            trans.append(list_line[0:3, 3])
            # 旋转矩阵转四元数
            r = Rotation.from_matrix(list_line[0:3, 0:3])
            quats.append(r.as_quat())

        # 最后得到两个numpu矩阵，dataset是存放所有真值的矩阵，groundtruth是存放xy真值的矩阵
        trans = np.array(trans)
        quats = np.array(quats)

        assert (trans.shape[0] == len(times))

        # 创建一个bag文件
        bag_out = rosbag.Bag(bagname, 'w')
        # 设置path消息的poses
        for i in range(len(trans)):
            pose = PoseStamped()
            pose.header.frame_id = 'odom'
            pose.header.stamp = rospy.Time(times[i])
            pose.header.seq = i
            pose.pose.position.x = trans[i][2]
            pose.pose.position.y = -trans[i][0]
            pose.pose.position.z = -trans[i][1]
            pose.pose.orientation.x = quats[i][2]
            pose.pose.orientation.y = -quats[i][0]
            pose.pose.orientation.z = -quats[i][1]
            pose.pose.orientation.w = quats[i][3]
            path.poses.append(pose)
            # 写入bag文件
            bag_out.write('/kitti/path', path, pose.header.stamp)

        # 关闭bag文件
        bag_out.close()


def write_bag_tum(filename, bagname, path, topic='/path'):
    trans = []
    quats = []
    times = []
    output = os.path.join(os.path.dirname(filename), bagname)
    print("output path: ", output)
    with open(filename) as f:
        list_file = f.readlines()
        # 将每一行数据转为数组
        for i in range(len(list_file)):
            list_line = list_file[i].split()
            # 将元素由字符串转为float
            list_line = list(map(float, list_line))
            times.append(list_line[0])

            list_line = np.array(list_line[1:])
            trans.append(list_line[0:3])
            quats.append(list_line[3:])

        # 最后得到两个numpu矩阵，dataset是存放所有真值的矩阵，groundtruth是存放xy真值的矩阵
        trans = np.array(trans)
        quats = np.array(quats)

        assert (trans.shape[0] == len(times))

        # 创建一个bag文件
        bag_out = rosbag.Bag(output, 'w', compression='bz2')
        # 设置path消息的poses
        for i in range(len(trans)):
            if i % 500 == 0:
                print(f"{i}/{len(trans)}")
            pose = PoseStamped()
            pose.header.frame_id = 'odom'
            pose.header.stamp = rospy.Time.from_sec(times[i])
            pose.header.seq = i
            pose.pose.position.x = trans[i][0]
            pose.pose.position.y = trans[i][1]
            pose.pose.position.z = trans[i][2]
            pose.pose.orientation.x = quats[i][0]
            pose.pose.orientation.y = quats[i][1]
            pose.pose.orientation.z = quats[i][2]
            pose.pose.orientation.w = quats[i][3]

            path.poses.append(pose)
            # 写入bag文件
            bag_out.write(topic, path, pose.header.stamp)

            if i > 3000:
                break
        # 关闭bag文件
        bag_out.close()


if __name__ == '__main__':
    rospy.init_node("gt2bag")

    # 创建一个path消息
    path = Path()
    path.header.frame_id = 'odom'
    path.header.stamp = rospy.Time.from_sec(0)

    # # kitti
    # gt_file = '/media/daybeha/Elements/SLAM_dataset/Kitti/sequences/poses/00.txt'
    # time_file = "/media/daybeha/Elements/SLAM_dataset/Kitti/sequences/00/times.txt"
    # bagname = "gt_00.bag"
    # times = get_times(time_file)
    # write_bag_kitti(gt_file, bagname, times, path)

    # # 4seasons
    # gt_file = '/media/daybeha/Elements/SLAM_dataset/4Seasons/recording_2020-03-24_17-36-22/result.txt'
    # bagname = "gt_office_loop_1_train.bag"
    # write_bag_tum(gt_file, bagname, path)

    # urban
    # 先用tools_py/trans_traj_gt_format.py把格式转为tum的 (不建议把urban39等这周几万帧的数据转为rosbag，太大了)
    gt_file = "/media/daybeha/LTFM2/dataset/urban/urban39-pankyo/global_pose_tum.txt"
    bagname = "gt_39.bag"
    write_bag_tum(gt_file, bagname, path, "/urban/path")

    print("success!")
